import logging
import json
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModel
import torch

# 配置日志
logging.basicConfig(
    level=logging.INFO,  # 设置日志级别为 INFO
    format="%(asctime)s - %(levelname)s - %(message)s",  # 设置日志格式
    handlers=[
        logging.FileHandler("server.log"),  # 将日志写入文件
        logging.StreamHandler()  # 同时将日志输出到控制台
    ]
)

app = FastAPI()

# 模型名称（也可以换成别的，比如：'BAAI/bge-small-zh'）
model_name = "/gemini/pretrain/bge-large-zh-v1.5"

# 加载 tokenizer 和 model
logging.info("正在加载模型和分词器...")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)
logging.info("模型和分词器加载完成。")

# 定义请求数据模型
class TextRequest(BaseModel):
    texts: list[str]

@app.post("/get_embeddings")
async def get_embeddings(request: TextRequest):
    try:
        texts = request.texts

        if not texts:
            raise HTTPException(status_code=400, detail="未提供有效的文本列表")

        # 编码输入文本
        logging.info("正在对输入文本进行编码...")
        inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
        logging.info("文本编码完成。")

        # 获取模型输出
        with torch.no_grad():
            logging.info("正在获取模型输出...")
            outputs = model(**inputs)
            embeddings = outputs.last_hidden_state[:, 0]  # 取 CLS token 的向量作为句向量

        # L2归一化（可选）
        logging.info("正在进行 L2 归一化...")
        embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)

        # 将嵌入向量转换为 Python 列表
        embedding_list = embeddings.tolist()

        # 构建 JSON 响应数据
        response_data = [
            {"text": text, "embedding": embedding}
            for text, embedding in zip(texts, embedding_list)
        ]

        logging.info("生成的 JSON 数据已准备好发送给客户端。")
        return response_data

    except Exception as e:
        logging.error(f"处理请求时发生错误: {str(e)}")
        raise HTTPException(status_code=500, detail="服务器内部错误")

if __name__ == "__main__":
    import uvicorn
    logging.info("启动服务端...")
    uvicorn.run(app, host="0.0.0.0", port=5000)  # 监听所有 IP 地址，端口为 5000